from .network_v2 import TactileDeltaEstimator
import torch
import numpy as np
import torch.nn as nn
import os

class TactileDeltaMeter:
    """
    Take a Tac3D frame as input and calculate the slip delta.
    The slip delta is an index that indicates how far the object is from macro slip.
    """

    def __init__(self, model: TactileDeltaEstimator, device):
        """
        Initialize the SlipDeltaMeter.

        """
        self.device = device
        self.model = model.to(self.device)
        self.mean = np.zeros((9), dtype=np.float32)
        self.std = np.ones((9), dtype=np.float32)
        self.online_training = False

    
    def frame_to_delta(self, frame):
        """
        Calculate the slip delta from a Tac3D frame.

        Args:
            frame (dict): The Tac3D frame.

        Returns:
            float: The slip delta.
        """
        # Preprocess the frame
        processed_data = self._frame_preprocess(frame)
        return self._get_network_output(processed_data)
        

    def data_to_delta(self, P,D,F):
        """
        Calculate the slip delta from Tac3D distributed data.

        Args:
            P (np.ndarray): The Tac3D distributed positions.
            D (np.ndarray): The Tac3D distributed displacements.
            F (np.ndarray): The Tac3D distributed forces.

        Returns:
            float: The slip delta.
        """
        # Preprocess the data
        processed_data = self._data_preprocess(P,D,F)
        return self._get_network_output(processed_data)

    
    def stacked_data_to_delta(self, data):
        """
        Calculate the slip delta from Tac3D stacked data.

        Args:
            data (np.ndarray): The Tac3D stacked data.

        Returns:
            float: The slip delta.
        """
        # Preprocess the data
        processed_data = self._stacked_data_preprocess(data)
        return self._get_network_output(processed_data)

    
    def _frame_preprocess(self, frame):
        """
        Preprocess the Tac3D frame.

        Args:
            frame (np.ndarray): The Tac3D frame.

        Returns:
            torch.tensor: The preprocessed Tac3D frame.
        """
        # Extract the Tac3D data
        P = frame["3D_Positions"]
        D = frame["3D_Displacements"]
        F = frame["3D_Forces"]
        
        return self._data_preprocess(P,D,F)
    
    def _data_preprocess(self, P,D,F):
        """
        Preprocess the Tac3D data.

        Args:
            P (np.ndarray): The Tac3D distributed positions.
            D (np.ndarray): The Tac3D distributed displacements.
            F (np.ndarray): The Tac3D distributed forces.

        Returns:
            torch.tensor: The preprocessed Tac3D data.
        """
        # Normalize the data
        data = np.concatenate((P, D, F), axis=1) # (400,3),(400,3),(400,3) -> (400,9)
        data = data[np.newaxis, ...]  # Add batch dimension
        return self._stacked_data_preprocess(data)
    
    def _stacked_data_preprocess(self, data):
        """
        Preprocess the Tac3D data.

        Args:
            data (np.ndarray): The Tac3D distributed data.

        Returns:
            torch.tensor: The preprocessed Tac3D data.
        """
        # Normalize the data
        data = (data - self.mean) / self.std

        # Convert to tensor
        data = torch.tensor(data, dtype=torch.float32)
        data = data.view(-1, 20, 20, 9).permute(0, 3, 1, 2).to(self.device) # (B,9,20,20)

        return data
    
    def _get_network_output(self, processed_data):
        """
        Get the network output.

        Args:
            processed_data (torch.tensor): The preprocessed Tac3D data.

        Returns:
            when online training: torch.tensor: The network output.
            when offline evaluation: np.ndarray: The network output.
        """
        if self.online_training:
            return  self.model(processed_data).squeeze()
        else:
            with torch.no_grad():
                return self.model(processed_data).cpu().numpy().squeeze()
    
    def back_propagate(self, output, target):
        """
        Backpropagate the loss.

        Args:
            output (torch.tensor): The network output.
            target (torch.tensor): The target slip delta.
        """
        assert self.online_training, "This method is only available in online training mode."
        assert type(output) == torch.Tensor, "Output must be a torch.Tensor."
        assert type(target) == torch.Tensor, "Target must be a torch.Tensor."
        if self.online_training:
            self.optimizer.zero_grad()
            loss = self.criterion(output, target)
            loss.backward()
            nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
            self.optimizer.step()
            self.scheduler.step(loss)
    
    def _get_res_dir(self):
        """
        Get the resource directory for the model and normalization parameters.

        Returns:
            str: The resource directory.
        """
        current_dir = os.path.dirname(os.path.abspath(__file__))
        res_idx = 0
        while True:
            res_dir = os.path.join(current_dir, f'res{res_idx}')
            if not (os.path.exists(res_dir) and os.path.exists(os.path.join(res_dir, "slip_delta_model.pth")) and os.path.exists(os.path.join(res_dir, "mean.npy")) and os.path.exists(os.path.join(res_dir, "std.npy"))):
                res_idx -= 1
                break
            res_idx += 1
        res_dir = os.path.join(current_dir, f'res{res_idx}')
        if not os.path.exists(res_dir):
            return None
        return res_dir
    
